package com.gitee.zyflzz.tool_set.huffmancoder;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * 哈夫曼编码器/解码器主类
 * 支持多线程处理和高并发操作
 *
 * @author Zhou Yifan
 */
public class HuffmanCoder {
    /**
     * 默认线程池大小 (核心线程数)
     */
    private final int DEFAULT_THREAD_COUNT = 2;
    /**
     * 线程池, 这里的是默认的
     */
    private ExecutorService executor = Executors.newFixedThreadPool(DEFAULT_THREAD_COUNT);
    private boolean defaultExecutor = true;

    public enum Case {
        ENCODE, DECODE
    }

    public HuffmanCoder() {
        this(null);
    }

    /**
     * 构造函数
     *
     * @param executor 必要的线程池
     */
    public HuffmanCoder(ExecutorService executor) {
        if (executor != null) {
            this.executor = executor;
            defaultExecutor = false;
        }
    }

    public ExecutorService getExecutor() {
        return executor;
    }

    /**
     * 指定任务类型进行编码或解码
     *
     * @param c         任务类型
     * @param filenames 文件名数组
     */
    public void run(Case c, String[] filenames) {

        List<File> files = Arrays.stream(filenames).map(File::new).collect(Collectors.toList());

        for (File file : files) {
            switch (c) {
                case ENCODE:
                    executor.submit(() -> {
                        try {
                            encode(file);
                        } catch (IOException e) {
                            // TODO: 可日志记录
                        }
                    });
                    break;
                case DECODE:
                    executor.submit(() -> {
                        try {
                            decode(file);
                        } catch (IOException e) {
                            // TODO: 可日志记录
                        }
                    });
                    break;
            }
        }

        if (defaultExecutor) {
            executor.shutdown();
            // try {
            //     executor.awaitTermination(1, TimeUnit.HOURS);
            // } catch (InterruptedException e) {
            //     // 可日志记录
            // }
        }
    }


    /**********************
     * 编码相关方法
     **********************/

    /**
     * 编码主方法
     *
     * @param file 需要编码的原始文件
     */
    public void encode(File file) throws IOException {
        // 步骤1：统计字节频率 (多线程优化)
        Map<Byte, Integer> frequency = computeFrequency(file);

        // 步骤2：构建哈夫曼树
        HuffmanTree tree = new HuffmanTree(frequency);

        // 步骤3：获取编码表
        Map<Byte, String> codeTable = tree.getCodeTable();

        // 步骤4：写入编码文件
        String encodedFileName = generateEncodedFileName(file);
        try (BitOutputStream bos = new BitOutputStream(Files.newOutputStream(Paths.get(encodedFileName)))) {
            // 写入原始文件名
            writeOriginalFileName(file, bos);
            // 写入频率表
            writeFrequencyTable(frequency, bos);
            // 写入编码后的数据
            writeEncodedData(file, codeTable, bos);
        }

        // TODO: 需要将原文件的文件格式写入?
    }

    /**
     * 多线程统计字节频率
     *
     * @param file 目标文件
     * @return 字节频率映射表
     */
    private Map<Byte, Integer> computeFrequency(File file) throws IOException {
        try (FileInputStream fis = new FileInputStream(file);
             FileChannel channel = fis.getChannel()) {
            // 获取文件总大小
            long fileSize = channel.size();
            // 设置分块大小 (1MB)
            int chunkSize = 1024 * 1024;

            // 创建专用线程池 (CPU核心数并行度)
            ExecutorService frequencyExecutor = Executors.newFixedThreadPool(
                    Runtime.getRuntime().availableProcessors()
            );

            // 存储各分块的Future结果
            List<Future<Map<Byte, Integer>>> futures = new ArrayList<>();

            // 分块处理文件
            long start = 0;
            while (start < fileSize) {
                long end = Math.min(start + chunkSize, fileSize);
                // 提交分块统计任务
                futures.add(frequencyExecutor.submit(
                        new FrequencyTask(channel, start, end)
                ));
                start = end;
            }

            // 合并所有分块的统计结果
            Map<Byte, Integer> frequency = new HashMap<>();
            for (Future<Map<Byte, Integer>> future : futures) {
                // 合并每个分块的统计结果
                frequency.putAll(future.get());
            }

            frequencyExecutor.shutdown();
            return frequency;
        } catch (InterruptedException | ExecutionException e) {
            throw new IOException("Error computing frequency", e);
        }
    }

    /**
     * 写入原始文件名到编码文件
     */
    private void writeOriginalFileName(File file, BitOutputStream bos) throws IOException {
        String fileName = file.getName();
        byte[] fileNameBytes = fileName.getBytes(StandardCharsets.UTF_8);
        bos.writeInt(fileNameBytes.length);    // 写入文件名长度 (4 字节)
        bos.write(fileNameBytes);              // 写入文件名内容
    }

    /**
     * 写入频率表到编码文件
     */
    private void writeFrequencyTable(Map<Byte, Integer> frequency,
                                     BitOutputStream bos) throws IOException {
        bos.writeInt(frequency.size());  // 写入表项数量 (4 字节)
        for (Map.Entry<Byte, Integer> entry : frequency.entrySet()) {
            bos.writeByte(entry.getKey());   // 写入字节 (8位)
            bos.writeInt(entry.getValue());  // 写入频率 (4 字节)
        }
    }

    /**
     * 写入编码数据
     */
    private void writeEncodedData(File file, Map<Byte, String> codeTable,
                                  BitOutputStream bos) throws IOException {
        try (FileInputStream fis = new FileInputStream(file);
             BufferedInputStream bis = new BufferedInputStream(fis)) {
            int b;
            while ((b = bis.read()) != -1) {  // 逐字节读取原始文件
                String code = codeTable.get((byte) b);  // 获取对应编码
                for (char c : code.toCharArray()) {     // 逐位写入
                    bos.writeBit(c == '1' ? 1 : 0);     // 1写入1, 0写入0
                }
            }
            bos.padLastByte();  // 填充最后一个字节的剩余位
        }
    }

    /**********************
     * 解码相关方法
     **********************/

    public void decode(File file) throws IOException {
        String decodedFileName = generateDecodedFileName(file);
        try (BitInputStream bis = new BitInputStream(Files.newInputStream(file.toPath()))) {
            // 读取原始文件名
            String originalFileName = readOriginalFileName(bis);
            // 文件格式还原的处理
            // 读取频率表
            Map<Byte, Integer> frequency = readFrequencyTable(bis);
            // 重建哈夫曼树
            HuffmanTree tree = new HuffmanTree(frequency);
            // 写入解码文件
            try (FileOutputStream fos = new FileOutputStream(decodedFileName);
                 BufferedOutputStream bos = new BufferedOutputStream(fos)) {
                decodeData(bis, tree.getRoot(), bos);
            }
        }
    }

    /**
     * 从编码文件读取原始文件名
     */
    private String readOriginalFileName(BitInputStream bis) throws IOException {
        int length = bis.readInt();          // 读取文件名长度
        byte[] fileNameBytes = new byte[length];
        for (int i = 0; i < length; i++) {
            fileNameBytes[i] = (byte) bis.readByte();  // 逐个字节读取文件名
        }
        return new String(fileNameBytes, StandardCharsets.UTF_8);
    }

    /**
     * 读取频率表
     */
    private Map<Byte, Integer> readFrequencyTable(BitInputStream bis) throws IOException {
        int size = bis.readInt();  // 读取表项数量
        Map<Byte, Integer> frequency = new HashMap<>();
        for (int i = 0; i < size; i++) {
            byte b = (byte) bis.readByte();  // 读取字节
            int freq = bis.readInt();        // 读取频率值
            frequency.put(b, freq);
        }
        return frequency;
    }

    /**
     * 解码数据核心逻辑: 路径法解码
     *
     * @param bis  位输入流
     * @param root 哈夫曼树根节点
     * @param bos  输出缓冲流
     */
    private void decodeData(BitInputStream bis, HuffmanTree.Node root,
                            BufferedOutputStream bos) throws IOException {
        HuffmanTree.Node current = root;
        int bit;
        while ((bit = bis.readBit()) != -1) {  // 逐位读取
            // 根据当前位选择路径
            current = (bit == 0) ? current.left : current.right;
            if (current.isLeaf()) {     // 到达叶子节点
                bos.write(current.b);   // 写入原始字节
                current = root;         // 重置到根节点
            }
        }
    }

    /**********************
     * 工具方法
     **********************/

    /**
     * 生成编码文件名 (原文件名_cod.huf)
     */
    private String generateEncodedFileName(File file) {
        String name = file.getName();
        int dotIndex = name.lastIndexOf('.');
        String base = (dotIndex == -1) ? name : name.substring(0, dotIndex);
        return file.getParent() + File.separator + base + "_cod.huf";
    }

    /**
     * 生成解码文件名 (原文件名_dec.原扩展名)
     */
    private String generateDecodedFileName(File file) {
        String name = file.getName();
        if (!name.endsWith("_cod.huf")) {
            throw new IllegalArgumentException("Invalid encoded file name");
        }
        String base = name.substring(0, name.length() - "_cod.huf".length());
        int dotIndex = base.lastIndexOf('.');
        String extension = (dotIndex == -1) ? "" : base.substring(dotIndex + 1);
        base = (dotIndex == -1) ? base : base.substring(0, dotIndex);
        return file.getParent() + File.separator + base + "_dec." + extension;
    }

    /**********************
     * 内部类定义
     **********************/

    /**
     * 频率统计任务 (实现Callable接口)
     */
    class FrequencyTask implements Callable<Map<Byte, Integer>> {
        /**
         * 文件通道
         */
        private final FileChannel channel;
        /**
         * 分块起始位置
         */
        private final long start;
        /**
         * 分块结束位置
         */
        private final long end;

        FrequencyTask(FileChannel channel, long start, long end) {
            this.channel = channel;
            this.start = start;
            this.end = end;
        }

        @Override
        public Map<Byte, Integer> call() throws Exception {
            Map<Byte, Integer> freq = new HashMap<>();
            // 分配与分块大小相同的缓冲区
            ByteBuffer buffer = ByteBuffer.allocate((int) (end - start));
            // 从指定位置读取文件内容
            channel.read(buffer, start);
            buffer.flip();  // 切换为读模式
            while (buffer.hasRemaining()) {
                byte b = buffer.get();  // 逐个字节读取
                freq.merge(b, 1, Integer::sum);  // 合并统计结果
            }
            return freq;
        }
    }

    /**
     * 哈夫曼树实现
     */
    class HuffmanTree {
        /**
         * 根节点
         */
        private final Node root;
        /**
         * 编码表, 键是存储的字节, 值是编码后的内容
         * 为了简便, 在这个阶段值就用String代替, 而不是真正的比特值
         */
        private final Map<Byte, String> codeTable;

        /**
         * 构造函数 (根据频率表构建树)
         */
        HuffmanTree(Map<Byte, Integer> frequency) {
            // 使用优先队列 (最小堆)构建哈夫曼树
            PriorityQueue<Node> pq = frequency.entrySet().stream()
                    .map(entry -> new Node(entry.getKey(), entry.getValue()))
                    .collect(Collectors.toCollection(PriorityQueue::new));

            // 合并节点直到只剩一个根节点
            while (pq.size() > 1) {
                Node left = pq.poll();
                Node right = pq.poll();
                // 创建新父节点 (频率为子节点之和)
                pq.add(new Node(null, left.freq + right.freq, left, right));
            }

            root = pq.poll();
            codeTable = new HashMap<>();
            buildCodeTable(root, "");  // 递归构建编码表
        }

        /**
         * 递归构建编码表 (前序遍历)
         */
        private void buildCodeTable(Node node, String code) {
            if (node == null) return;
            if (node.isLeaf()) {  // 叶子节点存储实际字节
                codeTable.put(node.b, code);
            } else {
                buildCodeTable(node.left, code + "0");  // 左子树添加0
                buildCodeTable(node.right, code + "1"); // 右子树添加1
            }
        }

        Map<Byte, String> getCodeTable() {
            return codeTable;
        }

        Node getRoot() {
            return root;
        }

        /**
         * 哈夫曼树节点定义
         */
        class Node implements Comparable<Node> {
            /**
             * 存储的字节 (叶子节点有效)
             */
            final Byte b;
            /**
             * 频率值
             */
            final int freq;
            /**
             * 左子节点
             */
            final Node left;
            /**
             * 右子节点
             */
            final Node right;

            /**
             * 构造叶子节点
             *
             * @param b    存储的字节 (叶子节点有效)
             * @param freq 频率值
             */
            Node(Byte b, int freq) {
                this(b, freq, null, null);
            }

            /**
             * 构造内部节点
             *
             * @param b     存储的字节 (叶子节点有效)
             * @param freq  频率值
             * @param left  左子节点
             * @param right 右子节点
             */
            Node(Byte b, int freq, Node left, Node right) {
                this.b = b;
                this.freq = freq;
                this.left = left;
                this.right = right;
            }

            /**
             * 判断是否是叶子节点
             */
            boolean isLeaf() {
                return left == null && right == null;
            }

            /**
             * 实现Comparable接口 (按频率排序)
             */
            @Override
            public int compareTo(Node other) {
                return this.freq - other.freq;
            }
        }
    }

    /**
     * 位输出流 (处理按位写入)
     */
    class BitOutputStream implements AutoCloseable {
        /**
         * 底层输出流
         */
        private final OutputStream out;
        /**
         * 当前缓冲的字节
         */
        private int currentByte;
        /**
         * 当前字节已使用的位数
         */
        private int bitsInCurrentByte;

        BitOutputStream(OutputStream out) {
            this.out = out;
        }

        /**
         * 写入单个位 (核心方法)
         */
        void writeBit(int bit) throws IOException {
            currentByte = (currentByte << 1) | bit;  // 左移, 并使用或运算添加新位
            bitsInCurrentByte++;
            if (bitsInCurrentByte == 8) {  // 满 8 位时写入 (因为计算机仅支持字节操作)
                out.write(currentByte);
                currentByte = 0;
                bitsInCurrentByte = 0;
            }
        }

        /**
         * 写入整数 (4 字节)
         */
        void writeInt(int value) throws IOException {
            for (int i = 24; i >= 0; i -= 8) {  // 高位优先
                writeByte((value >> i) & 0xFF);
            }
        }

        /**
         * 写入单字节 (8 位)
         */
        void writeByte(int value) throws IOException {
            for (int i = 7; i >= 0; i--) {  // 高位到低位
                writeBit((value >> i) & 1);
            }
        }

        /**
         * 写入字节数组
         */
        void write(byte[] bytes) throws IOException {
            for (byte b : bytes) {
                writeByte(b & 0xFF);
            }
        }

        /**
         * 填充最后一个字节的剩余位
         */
        void padLastByte() throws IOException {
            if (bitsInCurrentByte > 0) {
                currentByte <<= (8 - bitsInCurrentByte);  // 左移填充0
                out.write(currentByte);
                bitsInCurrentByte = 0;
            }
        }

        @Override
        public void close() throws IOException {
            padLastByte();  // 确保最后字节被写入
            out.close();
        }
    }

    /**
     * 位输入流 (处理按位读取)
     */
    class BitInputStream implements AutoCloseable {
        /**
         * 底层输入流
         */
        private final InputStream in;
        /**
         * 当前读取的字节
         */
        private int currentByte;
        /**
         * 当前字节剩余未读位数
         */
        private int bitsRemaining;

        BitInputStream(InputStream in) {
            this.in = in;
        }

        /**
         * 读取单个位 (核心方法)
         */
        int readBit() throws IOException {
            if (bitsRemaining == 0) {
                currentByte = in.read();  // 读取新字节
                if (currentByte == -1) return -1;  // 文件结束
                bitsRemaining = 8;        // 重置剩余位数
            }
            // 从高位到低位提取
            int bit = (currentByte >> (bitsRemaining - 1)) & 1;
            bitsRemaining--;
            return bit;
        }

        // 读取整数 (4 字节)
        int readInt() throws IOException {
            int value = 0;
            for (int i = 0; i < 4; i++) {
                value = (value << 8) | (readByte() & 0xFF);
            }
            return value;
        }

        // 读取单字节 (8 位)
        int readByte() throws IOException {
            int b = 0;
            for (int i = 0; i < 8; i++) {
                int bit = readBit();
                if (bit == -1) throw new EOFException();
                b = (b << 1) | bit;  // 左移并添加新位
            }
            return b;
        }

        @Override
        public void close() throws IOException {
            in.close();
        }
    }
}
